// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once
#include <set>
#include <string>
#include <vector>
#include <kllm/utility/all.h>
#include <turbo/flags/app.h>
#include <functional>

//
// CLI argument parsing
//
namespace kllm {

    struct AppArgs;
    using options_handler = std::function<void(KMParams &params, turbo::cli::App *app, const AppArgs &args)>;
    struct AppArgs {
        std::set<enum llama_example> examples = {LLAMA_EXAMPLE_COMMON};
        std::string args;
        std::string value_hint; // help text or example for arg value
        std::string value_hint_2; // for second arg value
        std::string env;
        std::string help;
        bool is_sparam = false; // is current arg a sampling param?
        options_handler app_handler;

        AppArgs(
                const std::string &args,
                const std::string &value_hint,
                const std::string &help,
                options_handler handler
        ) : args(args), value_hint(value_hint), help(help), app_handler(handler) {}

        AppArgs(
                const std::string &args,
                const std::string &help,
                options_handler handler
        ) : args(args), help(help), app_handler(handler)  {}

        // support 2 values for arg
        AppArgs(
                const std::string &args,
                const std::string &value_hint,
                const std::string &value_hint_2,
                const std::string &help,
                options_handler handler
        ) : args(args), value_hint(value_hint), value_hint_2(value_hint_2), help(help),  app_handler(handler) {}

        AppArgs &set_examples(std::initializer_list<enum llama_example> examples);

        AppArgs &set_env(const char *env);

        AppArgs &set_sparam();

        bool in_example(enum llama_example ex);

        bool get_value_from_env(std::string &output);

        bool has_value_from_env();

        std::string to_string();
    };

    struct ParamsContext {
        enum llama_example ex = LLAMA_EXAMPLE_COMMON;
        KMParams *params{nullptr};
        std::vector<AppArgs> options;

        static ParamsContext setup_app_context(turbo::cli::App *app, KMParams &params, llama_example ex);

        void add_option(AppArgs arg);
    };

    // function to be used by test-arg-parser
    void params_parser_init(ParamsContext&, KMParams &params, llama_example ex);
}  // namespace kllm
